import tensorflow as tf
import numpy as np
import random

class Cifar10_DataLoader:
    def __init__(self):
        (self.x_train, y_train), (self.x_test, y_test) = tf.keras.datasets.cifar10.load_data()
        self.y_train = y_train.reshape([-1])
        self.y_test = y_test.reshape([-1])

        self.dataSet_train = []
        self.dataSet_test = []
        for i in range(10):
            self.dataSet_train.append(np.where(self.y_train == i)[0])
            self.dataSet_test.append(np.where(self.y_test == i)[0])

        self.train_set_counts = len(self.dataSet_train[0])
        self.test_set_count = len(self.dataSet_test[0])

    def get_train_count(self):
        return self.train_set_counts

    def get_test_count(self):
        return self.test_set_count

    def get_train_batch(self, batch_size, label):
        rest = list(range(10))
        rest.remove(label)
        idx = []
        for _ in range(batch_size):
            if random.random()>0.5:
                idx.append(random.choice(self.dataSet_train[label]))
            else:
                idx.append(random.choice(self.dataSet_train[random.choice(rest)]))

        return self.x_train[idx], (self.y_train[idx]==label).astype(np.uint8)

    def get_test_batch(self, batch_size, label):
        rest = list(range(10))
        rest.remove(label)
        idx = []
        for _ in range(batch_size):
            if random.random() > 0.5:
                idx.append(random.choice(self.dataSet_test[label]))
            else:
                idx.append(random.choice(self.dataSet_test[random.choice(rest)]))

        return self.x_test[idx], (self.y_test[idx] == label).astype(np.uint8)

    def refresh(self):
        self.data_idx = 0

    def global_proprocessing(self, dataSet):
        pass

    def local_proprocessing(self, dataSet):
        pass

    def data_balance(self, data):
        from collections import Counter
        counter = Counter(data.reshape([-1]))
        print(counter)
